#include "bi_traversal.hpp"

namespace hms::swe
{

template<typename MeshType>
void computeBlocks(
	SolverThread<MeshType>& threadData,
	const IndexHelper& ih
){
	using B = BlockLocation;
	#define COMPUTE_BLOCK(LOC, BEG_X, BEG_Y, SIZE_X, SIZE_Y) \
	threadData.computeBlock( Block<B::LOC>{ih, BEG_X, BEG_Y, SIZE_X, SIZE_Y} )

	#define COMPUTE_CHECK_SIZE(LOC, BEG_X, BEG_Y, SIZE_X, SIZE_Y) \
	{ if (SIZE_X*SIZE_Y > 0) COMPUTE_BLOCK(LOC, BEG_X, BEG_Y, SIZE_X, SIZE_Y); }

	Index cellBegX, cellBegY;
	const Index
		& sizeX { ih.blocks.sizeX },
		& sizeY { ih.blocks.sizeY },
		& remSizeX { ih.blocks.remSizeX },
		& remSizeY { ih.blocks.remSizeY },
		& nol { ih.mesh.nOuterLayers };

	// std::stringstream s;
	// s
	// 	<< "Thread #" << omp_get_thread_num()
	// 	<< ", np=" << omp_get_num_threads()
	// 	<< "\n";
	// std::cout << s.str();

	/* inner blocks, full size */
	OMP_PRAGMA( for schedule(static) collapse(2) nowait )
	for (Index j=0; j<ih.blocks.nFullY; ++j ){
	for (Index i=0; i<ih.blocks.nFullX; ++i ){
		cellBegX = ih.beg.blocks + i*sizeX;
		cellBegY = ih.beg.blocks + j*sizeY;
		COMPUTE_BLOCK(inner, cellBegX, cellBegY, sizeX, sizeY);
	}}
	/* inner blocks, remainders */
	/* right */
	if (remSizeX > 0){
		OMP_PRAGMA( for schedule(dynamic) nowait )
		for (Index j=0; j<ih.blocks.nFullY; ++j ){
			cellBegY = ih.beg.blocks + j*sizeY;
			COMPUTE_BLOCK(inner, ih.beg.remX, cellBegY, remSizeX, sizeY);
		}
	}
	/* top */
	if (remSizeY > 0){
		OMP_PRAGMA( for schedule(dynamic) nowait )
		for (Index i=0; i<ih.blocks.nFullX; ++i ){
			cellBegX = ih.beg.blocks + i*sizeX;
			COMPUTE_BLOCK(inner, cellBegX, ih.beg.remY, sizeX, remSizeY);
		}
	}

	/* outer blocks, bottom and top */
	OMP_PRAGMA( for schedule(dynamic) nowait )
	for ( Index i=0; i<ih.blocks.nFullX; ++i ){
		cellBegX = ih.beg.blocks + i*sizeX;
		COMPUTE_BLOCK(bottom, cellBegX, ih.beg.mesh, sizeX, nol);
		/* top block */
		COMPUTE_BLOCK(top, cellBegX, ih.beg.outerY, sizeX, nol);
	}
	/* outer blocks, left and right */
	OMP_PRAGMA( for schedule(dynamic) nowait )
	for ( Index j=0; j<ih.blocks.nFullY; ++j ){
		cellBegY = ih.beg.blocks + j*sizeY;
		COMPUTE_BLOCK(left , ih.beg.mesh  , cellBegY, nol, sizeY);
		/* right */
		COMPUTE_BLOCK(right, ih.beg.outerX, cellBegY, nol, sizeY);
	}

	/* remainder blocks and corners */
	/* removed nowait clause, because it caused issues with the swap of old and
	 * new fields. And generally, after all blocks have been computed, a
	 * barrier does make sense. */
	OMP_PRAGMA( sections ) // nowait
	{
		/* remainder corner, full muscl */
		OMP_PRAGMA( section )
		COMPUTE_CHECK_SIZE(inner, ih.beg.remX, ih.beg.remY, remSizeX, remSizeY)

		/* remainder bottom, mixed */
		OMP_PRAGMA( section )
		COMPUTE_CHECK_SIZE(bottom, ih.beg.remX, ih.beg.mesh  , remSizeX, nol)
		/* remainder top, mixed */
		OMP_PRAGMA( section )
		COMPUTE_CHECK_SIZE(top   , ih.beg.remX, ih.beg.outerY, remSizeX, nol)

		/* remainder left, mixed */
		OMP_PRAGMA( section )
		COMPUTE_CHECK_SIZE(left , ih.beg.mesh  , ih.beg.remY, nol, remSizeY)
		/* remainder right, mixed */
		OMP_PRAGMA( section )
		COMPUTE_CHECK_SIZE(right, ih.beg.outerX, ih.beg.remY, nol, remSizeY)

		/* corners, mixed */
		OMP_PRAGMA( section )
		{
			COMPUTE_BLOCK(bottomLeft , ih.beg.mesh  , ih.beg.mesh  , nol, nol);
			COMPUTE_BLOCK(bottomRight, ih.beg.outerX, ih.beg.mesh  , nol, nol);
			COMPUTE_BLOCK(topLeft    , ih.beg.mesh  , ih.beg.outerY, nol, nol);
			COMPUTE_BLOCK(topRight   , ih.beg.outerX, ih.beg.outerY, nol, nol);
		}
	}

	#undef COMPUTE_BLOCK
	#undef COMPUTE_CHECK_SIZE
}

template void computeBlocks( SolverThread<UniMesh>&, const IndexHelper& );
template void computeBlocks( SolverThread<RectMesh>&, const IndexHelper& );


} // namespace hms::swe
